
#include<iostream>
using namespace std;


#include "PacketSnifferWorkerThread.h"
#include "PacketSniffer.h"

PacketSnifferWorkerThread::PacketSnifferWorkerThread
( 
    QWaitCondition* notify,
    pcap_if_t* device,
    const int& timeout, // Milliseconds.
    const int& mode, // Promiscuous.
    const QString& filter, // Filter expression.
    const bool& tcpOnly,
    const int& packetPortion )
    :
    NULL_HDRLEN( 4 ),
    PPP_HDRLEN( 4 ),
    SLL_HDRLEN( 16 ),
    initSuccess_( false ),
    notify_( notify ),
    device_( device ),
    timeout_( timeout ),
    mode_( mode ),
    tcpOnly_( tcpOnly ),
    packetPortion_( packetPortion )
{
    // Return code for things.
    int rc;

    // Prevent non-IP packets from being delivered.
    QString filterAndIP = QString("(ip) and (") + filter + QString(")");

    // Allocate chars for the filter and copy it over.
    filterStr_ = new char[ filterAndIP.length()+1 ];
    memcpy( filterStr_, filterAndIP.ascii(), filterAndIP.length()+1 );

    // Try to do an open.
    adhandle_ = pcap_open_live( device_->name,
				packetPortion_,
				mode_,
				timeout_,
				errbuf );
    // Check to see if it failed or succeeded.
    if( adhandle_ == NULL )
    {
	// Failed.
        cerr << "Unable to open adapter.  "
	     << device_->name << " is not supported." << endl;
	initSuccess_ = false;
	return;
    }

    // Get the link type.
    dataLinkType_ = pcap_datalink( adhandle_ );

    // If the filter is set, then set one.
    if( filter != "" )
    {
	// Get the netmask because the compiler filter sometimes needs it.
	if( device_->addresses != 0 )
	{
	    /* Retrieve the mask of the first address of the interface */
	    netmask_ = ((struct sockaddr_in *)(device_->addresses->netmask))->
		sin_addr.S_un.S_addr;
	}
	else
	{
	    /* If the interface is without addresses we suppose to be in a
	     * C class network */
	    netmask_=0xffffff; 
	}
    

	// Compile the filter
	rc = pcap_compile( adhandle_, &fcode_, filterStr_, 1, netmask_ );
	if( rc < 0 ) 
	{
	    cerr << "Unable to compile the packet filter. Check the syntax." 
		 << endl;
	    initSuccess_ = false;
	    return;
	}
    
	//set the filter
	rc = pcap_setfilter( adhandle_, &fcode_ );
	if( rc < 0 )
	{
	    cerr << "ERROR:  Error setting the filter." << endl;
	    initSuccess_ = false;
	    return;
	}
    }

//     debugworker.open(string(string("debugworker") + string(QString::number((long)(this)).ascii()) + string(".txt")).c_str(), ios::out | ios::app);
//     if( !debugworker.is_open() )
//     {
// 	exit();
//     }
//     struct ip x;
//     debugworker << "DEBUG:  printing sizes." << endl;
//     debugworker << "DEBUG:  size of ip: " << sizeof(x) << endl;    
// //     debugworker << "DEBUG:  size of ip_hl: " << sizeof((x.ip_hl)) << endl;
// //     debugworker << "DEBUG:  size of ip_v: " << sizeof((x.ip_v)) << endl;
// //    debugworker << "DEBUG:  size of ip_tos: " << sizeof((x.ip_tos)) << endl;
//     debugworker << "DEBUG:  size of ip_len: " << sizeof((x.ip_len)) << endl;
//     debugworker << "DEBUG:  size of ip_id: " << sizeof((x.ip_id)) << endl;
//     debugworker << "DEBUG:  size of ip_off: " << sizeof((x.ip_off)) << endl;
//     debugworker << "DEBUG:  size of ip_ttl: " << sizeof((x.ip_ttl)) << endl;
//     debugworker << "DEBUG:  size of ip_p: " << sizeof((x.ip_p)) << endl;
//     debugworker << "DEBUG:  size of ip_sum: " << sizeof((x.ip_sum)) << endl;
//     debugworker << "DEBUG:  size of ip_src: " << sizeof((x.ip_src)) << endl;
//     debugworker << "DEBUG:  size of ip_dst: " << sizeof((x.ip_dst)) << endl;

//     debugworker << "DEBUG:  size of ip: " << (long)(&x) << endl;    
// //     debugworker << "DEBUG:  size of ip_hl: " << (long)(&(x.ip_hl)) << endl;
// //     debugworker << "DEBUG:  size of ip_v: " << (long)(&(x.ip_v)) << endl;
// //    debugworker << "DEBUG:  size of ip_tos: " << (long)(&(x.ip_tos)) << endl;
//     debugworker << "DEBUG:  size of ip_len: " << (long)(&(x.ip_len)) << endl;
//     debugworker << "DEBUG:  size of ip_id: " << (long)(&(x.ip_id)) << endl;
//     debugworker << "DEBUG:  size of ip_off: " << (long)(&(x.ip_off)) << endl;
//     debugworker << "DEBUG:  size of ip_ttl: " << (long)(&(x.ip_ttl)) << endl;
//     debugworker << "DEBUG:  size of ip_p: " << (long)(&(x.ip_p)) << endl;
//     debugworker << "DEBUG:  size of ip_sum: " << (long)(&(x.ip_sum)) << endl;
//     debugworker << "DEBUG:  size of ip_src: " << (long)(&(x.ip_src)) << endl;
//     debugworker << "DEBUG:  size of ip_dst: " << (long)(&(x.ip_dst)) << endl;


    
    // Success.
    initSuccess_ = true;
    
}

PacketSnifferWorkerThread::~PacketSnifferWorkerThread()
{
    delete [] filterStr_;
    filterStr_ = 0;

//     debugworker.close();
}

QByteArray PacketSnifferWorkerThread::getByteArray()
{
    if( initSuccess_ )
    {
	QByteArray rv = bytes_.copy();
	bytes_.resize(0);
	return rv;
    }
    else
    {
	bytes_ = QByteArray();
	return bytes_;
    }
}


bool PacketSnifferWorkerThread::initSuccess() const
{
    return initSuccess_;
}

bool PacketSnifferWorkerThread::packetReady() const
{
    // Checks the QByteArray if it's empty or not.
    if( ! bytes_.isEmpty() ) 
    {
	return true;
    }
    else 
    {
	return false;
    }
}


void PacketSnifferWorkerThread::run()
{
    // Make sure we init-ed okay.
    if( !initSuccess_ )
    {
	return;
    }

    
    struct pcap_pkthdr* pkt_header;
    u_char* pkt_data;
    
    // Initialize the return code.
    int rc = 0;

// Possible return codes.
//     * 1 if the packet has been read without problems
//     * 0 if the timeout set with pcap_open_live() has elapsed. In
//     *       this case pkt_header and pkt_data don't point to a valid
//     *       packet
//     * -1 if an error occurred
//     * -2 if EOF was reached reading from an offline capture

    // While we haven't got any valid data and it's not an error, try again.
    while( rc == 0 ) 
    {
	rc = pcap_next_ex( adhandle_, 
			   &pkt_header,
			   &pkt_data );
    }

    switch( rc )
    {
    case 1: 
    {
	// Packet ready!

	// Strip any extra headers before the IP headers.
	switch (dataLinkType_)
	{
	case DLT_NULL:
	    bytes_ = handle_dl_null( pkt_header, pkt_data );
            // 	    debugworker << "PACKET type is NULL" << endl;
	    break;
	case DLT_RAW:
	    bytes_ = handle_dl_raw( pkt_header, pkt_data );
            // 	    debugworker << "PACKET type is RAW" << endl;
	    break;
	case DLT_EN10MB:
	    bytes_ = handle_dl_ethernet( pkt_header, pkt_data );
            // 	    debugworker << "PACKET type is ethernet" << endl;
	    break;
	case DLT_IEEE802:
	    bytes_ = handle_dl_ethernet( pkt_header, pkt_data );
            // 	    debugworker << "PACKET type is ethernet - 802" << endl;
	    break;
	case DLT_PPP:
	    bytes_ = handle_dl_ppp( pkt_header, pkt_data );
	    break;
	case DLT_LINUX_SLL:
	    bytes_ = handle_dl_linux_sll( pkt_header, pkt_data );
	    break;
	default:
	    cerr << "WARNING:  Unknown data link type." << endl;
	    bytes_.resize(0);
            // 	    debugworker << "PACKET type unknown" << endl;
            // 	    debugworker.flush();
	    return;
	}


	

	// Turn it into only TCP data if necessary.
	if( (tcpOnly_ == true) && (bytes_.isEmpty() == false) )
	{
	    bytes_ = stripTCP( bytes_ );
	}
	
	// Wake everyone on the wait condition.
	notify_->wakeAll();
	// Complete the thread.
	break;
    }
    case -1: 
    {
	// Error occurred.
	cerr << "ERROR:  " << errbuf << endl;
	// Complete the thread.
	break;
    }
    case -2:
    {
	// EOF.
	cerr << "ERROR:  Reached end of file." << endl;
	// Complete the thread.
	break;
    }
    default:
    {
	// Unknown return code.
	cerr << "ERROR:  Unknown return code (" << rc << ")" << endl;
    }
    }  // End of     switch( rc )

} // End of    run().




int PacketSnifferWorkerThread::getDataLinkType() const
{
    return dataLinkType_;
}


/**
 * This handles a null packet header type.  The DLT_NULL packet header
 * is 4 bytes long (see NULL_HDRLEN in header file).  It contains a
 * network order 32-bit integer that specifies the family,
 * e.g. AF_INET.  DLT_NULL is used by the localhost interface.
 *
 * @return Returns the QByteArray containing the IP packet for this data link
 * type.  
 */
QByteArray 
PacketSnifferWorkerThread::handle_dl_null( const struct pcap_pkthdr* header, 
					   const unsigned char* p )
{
    unsigned int caplen = header->caplen;
    unsigned int len = header->len;
    unsigned int family;
    
    if( caplen != len )
    {
        // 	cout << "WARNING:  Only captured " << caplen << " bytes of "
        // 	     << len << " null frame."  << endl;
    }
    
    if( caplen < NULL_HDRLEN )
    {
	//      cout << "WARNING:  Received incomplete null frame." << endl;
	return QByteArray();
    }
    
    // Check for AF_INET family.
    memcpy((char *)&family, (char *)p, sizeof(family));
    family = ntohl(family);
    if( family != AF_INET )
    {
	//      cout << "WARNING:  Received non- AF_INET Null frame (type " 
        //	     << family << ")." << endl;
	return QByteArray();
    }

    // Nicely named helper variables.
    unsigned int actualIPPacketSize = caplen - NULL_HDRLEN;
    const unsigned char* actualIPPacket = p + NULL_HDRLEN;

    // Check for valid IP packet.
    if( ! checkValidIPPacket( actualIPPacket, actualIPPacketSize ) )
    {
	//      cout << "WARNING:  Packet was found to be invalid." << endl;
	return QByteArray();
    }
    
    // Create the return value.
    QByteArray rv( actualIPPacketSize );

    // Copy over the memory.
    memcpy( rv.data(), actualIPPacket, actualIPPacketSize );

    // Return the IP packet.
    return rv;
}


/**
 * This handles a Raw IP Packet.  No encapsulation or link-layer headers.
 * Used for PPP connections under some OSs including Linux and IRIX.
 *
 * @return Returns the QByteArray containing the IP packet for this data link
 * type.
 */
QByteArray 
PacketSnifferWorkerThread::handle_dl_raw( const struct pcap_pkthdr* header, 
					  const unsigned char* p )
{
    unsigned int caplen = header->caplen;
    unsigned int len = header->len;
    
    if( len != caplen )
    {
        // 	cout << "WARNING:  Only captured " << caplen << " bytes of "
        // 	     << len << "-byte raw frame." << endl;
    }

    // Nicely named helper variables.
    unsigned int actualIPPacketSize = caplen;
    const unsigned char* actualIPPacket = p;

    // Check for valid IP packet.
    if( ! checkValidIPPacket( actualIPPacket, actualIPPacketSize ) )
    {
	// cout << "WARNING:  Packet was found to be invalid." << endl;
	return QByteArray();
    }

    // Create the return value.
    QByteArray rv( actualIPPacketSize );

    // Copy over the memory.
    memcpy( rv.data(), actualIPPacket, actualIPPacketSize );

    // Return the IP packet.
    return rv;
}


/**
 * This handles a Ethernet datalink packet.  This is used by all 10
 * and 100 mbit/sec ethernet.  We are given the entire ethernet header
 * so we check to make sure it's marked as being IP.
 *
 * @return Returns the QByteArray containing the IP packet for this
 * data link type.
 */
QByteArray 
PacketSnifferWorkerThread::handle_dl_ethernet( const struct pcap_pkthdr* header,
					       const unsigned char* p )
{
    unsigned int caplen = header->caplen;
    unsigned int len = header->len;
    struct ether_header *eth_header = (struct ether_header *)p;
    


    if( len != caplen )
    {
// 	debugworker << "WARNING:  Only captured " << caplen << " bytes of " 
// 	     << len << "-byte ethernet frame." << endl;
// 	debugworker.flush();
    }
    
    if( caplen < sizeof(struct ether_header) )
    {
// 	debugworker << "WARNING:  Received incomplete ethernet frame." << endl;
// 	debugworker.flush();
	return QByteArray();
    }
    
    // We're only expected IP datagrams, nothing else.
    if( ntohs( eth_header->ether_type ) != ETHERTYPE_IP )
    {
// 	debugworker << "WARNING:  Received ethernet frame with unknown type "
// 	     << ntohs( eth_header->ether_type );
// 	debugworker.flush();
	
	return QByteArray();
    }
    
    // Nicely named helper variables.
    unsigned int actualIPPacketSize = caplen - sizeof(struct ether_header);
    const unsigned char* actualIPPacket = p + sizeof(struct ether_header);

    // Check for valid IP packet.
    if( ! checkValidIPPacket( actualIPPacket, actualIPPacketSize ) )
    {
// 	debugworker << "WARNING:  Packet was found to be invalid." << endl;
// 	debugworker.flush();
	return QByteArray();
    }

    // Create the return value.
    QByteArray rv( actualIPPacketSize );

    // Copy over the memory.
    memcpy( rv.data(), actualIPPacket, actualIPPacketSize );

//     debugworker << "leave handle_dl_ethernet()" << endl;
    
    // Return the IP packet.
    return rv;
}


/**
 * This handles the PPP datalink packet.  The DLT_PPP packet header is
 * 4 bytes long.  We just move past it without parsing it.  It is used
 * for PPP on some OSs (DLT_RAW is used by others; see
 * handle_dl_raw()).
 * 
 * @return Returns the QByteArray containing the IP packet for this data link type.
 */
QByteArray 
PacketSnifferWorkerThread::handle_dl_ppp( const struct pcap_pkthdr* header, 
					  const unsigned char* p )
{
    unsigned int caplen = header->caplen;
    unsigned int len = header->len;

    if( len != caplen )
    {
	// cout << "WARNING:  Only captured " << caplen << " bytes of "
	//      << len << "-byte PPP frame." << endl;
    }
    
    if( caplen < PPP_HDRLEN )
    {
	// cout << "WARNING:  Received incomplete PPP frame" << endl;
	return QByteArray();
    }
    
    // Nicely named helper variables.
    unsigned int actualIPPacketSize = caplen - PPP_HDRLEN;
    const unsigned char* actualIPPacket = p + PPP_HDRLEN;

    // Check for valid IP packet.
    if( ! checkValidIPPacket( actualIPPacket, actualIPPacketSize ) )
    {
	// cout << "WARNING:  Packet was found to be invalid." << endl;
	return QByteArray();
    }

    // Create the return value.
    QByteArray rv( actualIPPacketSize );

    // Copy over the memory.
    memcpy( rv.data(), actualIPPacket, actualIPPacketSize );

    // Return the IP packet.
    return rv;
}


/**
 * Returns the QByteArray containing the IP packet for this data link type.
 */
QByteArray 
PacketSnifferWorkerThread::handle_dl_linux_sll( const struct pcap_pkthdr* header, 
						const unsigned char* p )
{
    unsigned int caplen = header->caplen;
    unsigned int len = header->len;
   
    if( len != caplen )
    {
	// cout << "WARNING:  Only captured " << caplen << " bytes of "
	//      << len << "-byte Linux cooked frame." << endl;
    }
    
    if( caplen < SLL_HDRLEN )
    {
	// cout << "WARNING:  Received incomplete Linux cooked frame" << endl;
	return QByteArray();
    }
    
    // Nicely named helper variables.
    unsigned int actualIPPacketSize = caplen - SLL_HDRLEN;
    const unsigned char* actualIPPacket = p + SLL_HDRLEN;


    // Check for valid IP packet.
    if( ! checkValidIPPacket( actualIPPacket, actualIPPacketSize ) )
    {
	// cout << "WARNING:  Packet was found to be invalid." << endl;
	return QByteArray();
    }

    // Create the return value.
    QByteArray rv( actualIPPacketSize );

    // Copy over the memory.
    memcpy( rv.data(), actualIPPacket, actualIPPacketSize );

    // Return the IP packet.
    return rv;
    
}

/**
 * Returns if it found the IP packet to be valid as far as we can understand.
 */
bool 
PacketSnifferWorkerThread::checkValidIPPacket( const unsigned char* data, 
					       const unsigned int& caplen )
{
//     debugworker << "DEBUG:  enter checkValidIPPacket()" << endl;
    
    const struct ip* ip_header = (struct ip *)data;
    unsigned int ip_header_len;
    unsigned int ip_total_len;
    
//     debugworker << "DEBUG:  caplen == " << caplen << endl
// 		<< "DEBUG:  sizeof(struct ip) == " << sizeof(struct ip) << endl;
//     debugworker.flush();
    
    
    // Make sure the packet is at least as long as the min IP header.
    if( caplen < sizeof(struct ip) )
    {
// 	debugworker << "WARNING:  Received truncated IP datagram!" << endl;
// 	debugworker.flush();
	return false;
    }
    
    // Do a quick check for TCP protocol if tcpOnly_ is true.
    if( (tcpOnly_ == true) && (ip_header->ip_p != IPPROTO_TCP) ) 
    {
// 	debugworker << "WARNING: User specified tcpOnly but we received a "
// 		    << "non TCP packet!" << endl;
// 	debugworker.flush();
	
	return false;
    }
//     debugworker << "DEBUG:  version is:  " << ( ip_header->ip_hl_v_tos >> 4 ) << endl;


    // Check and make sure we got everything.  Note: We must use
    // ip_total_len after this, because we may have captured bytes
    // beyond the end of the packet (e.g. Ethernet padding).
//     debugworker << "DEBUG:  bytes of ip_total_len:  ";
//     unsigned char* loc = (unsigned char*)&(ip_header->ip_len);
    
//     debugworker << print_hex_form( *loc );
//     debugworker << print_hex_form( *(loc + 1) );
//     debugworker << print_hex_form( *(loc + 2) );
//     debugworker << print_hex_form( *(loc + 3) );
//     debugworker << "*******" << endl;

    
    ip_total_len = ntohs(ip_header->ip_len);

//     debugworker << "DEBUG:  ip_total_len == " << ip_total_len << endl;
//     debugworker.flush();
    
    if( caplen < ip_total_len )
    {
// 	debugworker << "WARNING:  Captured only " << caplen << " bytes of "
// 	     << ip_total_len << "-byte IP datagram." << endl;
// 	debugworker.flush();
	
    }

//     // DEBUG:
//     QByteArray debugin( caplen );
//     memcpy(debugin.data(), data, caplen );
//     print_qbytearray( debugin );

    // Throw away everything but fragment 0; I dont' know how to do
    // fragment reassembly.
    if( ntohs(ip_header->ip_off) & 0x1fff )
    {
// 	debugworker << "WARNING:  Throwing away IP fragment." << endl;
// 	debugworker.flush();
	
	return false;
    }
    
    
    // Figure out where the IP header ends.
    ip_header_len = ( (ip_header->ip_hl_v_tos & 0x000f ) ) * 4;
    
//     debugworker << "DEBUG:  ip_header_len == " << ip_header_len << endl;

    // Make sure there's some data.
    if( ip_header_len > ip_total_len )
    {
// 	debugworker << "WARNING:  Received truncated IP datagram!" << endl;
// 	debugworker.flush();
	return false;
    }
    
//     debugworker << "DEBUG:  PASSED IP validness test" << endl;
//     debugworker.flush();
    
    // It passed our tests.
    return true;
}



/**
 * Takes a IP packet, strips out the IP and TCP headers, and returns
 * just the TCP data part.
 */
QByteArray 
PacketSnifferWorkerThread::stripTCP( const QByteArray& IPpacket )
{
    const struct ip* ip_header = (struct ip *)IPpacket.data();
    unsigned int ip_header_len = ( ip_header->ip_hl_v_tos & 0x000f ) * 4;
    unsigned int ip_total_len  = ntohs(ip_header->ip_len);
    unsigned int tcp_total_len = ip_total_len - ip_header_len;
    unsigned int tcp_header_len;
    
    QByteArray TCPpacket( tcp_total_len );
    memcpy( TCPpacket.data(), IPpacket.data() + ip_header_len, tcp_total_len );
    
    struct tcphdr *tcp_header = (struct tcphdr *)TCPpacket.data();
    
    if( tcp_total_len < sizeof(struct tcphdr) )
    {
	// cout << "WARNING:  Received truncated TCP segment!" << endl;
	return QByteArray();
    }
    
    // Calculate the total length of the TCP header, including options.
    tcp_header_len = ( tcp_header->th_x2_off >> 4 ) * 4;
    
    // Return if this packet doesn't have any data (e.g., an ACK).
    if( tcp_total_len <= tcp_header_len )
    {
	// cout << "INFO:  Got a TCP segment without any data." << endl;
	return QByteArray();
    }

    unsigned int dataSize = tcp_total_len - tcp_header_len;
    QByteArray data( dataSize );
    memcpy( data.data(), TCPpacket.data() + tcp_header_len, dataSize );

    return data;
}

// string PacketSnifferWorkerThread::print_hex_form(unsigned char x)
// {
//       string rv = "0x";

//   if((x/16) == 10)
//     rv += "A";
//   else if((x/16) == 11)
//     rv += "B";
//   else if((x/16) == 12)
//     rv += "C";
//   else if((x/16) == 13)
//     rv += "D";
//   else if((x/16) == 14)
//     rv += "E";
//   else if((x/16) == 15)
//     rv += "F";
//   else
//     rv += char(48 + (x/16));

//   if((x%16) == 10)
//     rv += "A";
//   else if((x%16) == 11)
//     rv += "B";
//   else if((x%16) == 12)
//     rv += "C";
//   else if((x%16) == 13)
//     rv += "D";
//   else if((x%16) == 14)
//     rv += "E";
//   else if((x%16) == 15)
//     rv += "F";
//   else
//     rv += char(48 + (x%16));
  
//   rv += " ";
//   return rv;
// }

// void PacketSnifferWorkerThread::print_qbytearray( const QByteArray& arr )
// {
//     debugworker << "##############################################################################" << endl << "# DUMP PACKET " << endl;

//     if( arr.size() == 0 )
//     {
// 	debugworker << "OMFG BLANK!" << endl;
//     }
    
    
//     for( int j = 0; j < arr.size(); j++ )
//     {
// 	debugworker << print_hex_form(arr.data()[j]);
//     }
//     debugworker << endl;
//     debugworker.flush();
// }


